/*
 * Copyright (c) 2010 Network Security Lab (NSL), Department of Electrical
 *                    Engineering, University of Washington
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation;
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 * Author: He Wu mdzz@uw.edu
 */

/**
 * @file packet_handler.c
 */

#include "packet_handler.h"
#include "command.h"
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <assert.h>

/******************************************************************************/

/**
 * @brief Worker threads that process packets from packet queue.
 * @param q Pointer to task queue.
 */
static void worker(task_q *tq);

/******************************************************************************/

void packet_q_append(packet_q *q, packet *pkt)
{
  assert(q && pkt);

  pthread_mutex_lock(&(q->lock));

  // empty queue
  if (q->head == NULL) {
    q->head = malloc(sizeof(packet_q_node));
    q->head->pkt = pkt;
    q->head->next = NULL;
    q->tail = q->head;
    pthread_mutex_unlock(&(q->lock));
    return;
  }

  // append new element to the end of queue
  assert(q->tail);
  q->tail->next = malloc(sizeof(packet_q_node));
  q->tail->next->pkt = pkt;
  q->tail->next->next = NULL;
  q->tail = q->tail->next;

  pthread_mutex_unlock(&(q->lock));
}


packet *packet_q_remove(packet_q *q)
{
  assert(q);
  packet *pkt;
  packet_q_node *oldHead;

  pthread_mutex_lock(&(q->lock));
  // empty queue
  if (q->head == NULL) {
    pthread_mutex_unlock(&(q->lock));
    return NULL;
  }
  // remove head
  pkt = q->head->pkt;
  oldHead = q->head;
  q->head = q->head->next;
  // check if queue is empty
  if (q->head == NULL)
    q->tail = NULL; // set empty tail
  pthread_mutex_unlock(&(q->lock));

  free(oldHead);

  return pkt;
}


int packet_q_size(packet_q *q)
{
  assert(q);
  int i = 0;
  packet_q_node *cur;
  pthread_mutex_lock(&(q->lock));
  cur = q->head;
  while (cur) {
    cur = cur->next;
    i++;
  }
  pthread_mutex_unlock(&(q->lock));
  return i;
}


void packet_q_init(packet_q *q)
{
  assert(q);
  q->head = NULL;
  q->tail = NULL;
  pthread_mutex_init(&(q->lock), NULL);
}


void packet_q_destroy(packet_q *q)
{
  assert(q);

  pthread_mutex_lock(&(q->lock));
  packet_q_node *cur = q->head;
  packet_q_node *tmp = cur;
  while (cur) {
    tmp = cur;
    cur = cur->next;
    free(tmp->pkt);
    free(tmp);
  }
  q->head = NULL;
  q->tail = NULL;
  pthread_mutex_unlock(&(q->lock));

  pthread_mutex_destroy(&(q->lock));

  free(q);
}


int task_queue_init(task_q *tq, int num_of_workers)
{
  assert(tq);
  int i;

  printf("Creating task queue!\n");

  tq->q = malloc(sizeof(packet_q));
  packet_q_init(tq->q);
  tq->destroy = 0;
  tq->num_of_workers = num_of_workers;
  pthread_mutex_init(&(tq->lock), NULL);
  pthread_cond_init(&(tq->cv), NULL);

  if (num_of_workers == 0)
    return 0;     // single threaded

  tq->workers = malloc(sizeof(pthread_t *) * num_of_workers);
  if (tq->workers == NULL)
    return -1;    // allocation failed

  // start worker threads
  for (i = 0; i < num_of_workers; i++) {
    if (pthread_create(&(tq->workers[i]), NULL, (void *)&worker, (void *)tq))
      return -1;  // failed to create thread
  }

  printf("Done creating task queue!\n");

  return 0;
}

void destroy_task_queue(task_q *tq)
{
  assert(tq);
  int i;

  // set flag
  tq->destroy = 1;
  // wake up everyone
  pthread_cond_broadcast(&(tq->cv));
  // wait for all workers to finish
  for (i = 0; i < tq->num_of_workers; i++)
    pthread_join(tq->workers[i], NULL);
  // destroy packet queue
  packet_q_destroy(tq->q);
  // clean up the rest
  pthread_cond_destroy(&(tq->cv));
  if (tq->num_of_workers != 0)
    free(tq->workers);
  free(tq);
}

void add_task(task_q *tq, packet *pkt)
{
  assert (tq && pkt);

  if (tq->num_of_workers == 0) {
    pkt_processor(pkt);
    free(pkt);
  }

  // add to packet queue
  packet_q_append(tq->q, pkt);
  // signal workers
  pthread_cond_signal(&(tq->cv));
}

/******************************************************************************/


static void worker(task_q *tq)
{
  packet *pkt;

  printf("Worker started!\n");

  // worker runs forever
  while (1) {
    pthread_mutex_lock(&(tq->lock));

    // check if there's any packet
    while (packet_q_size(tq->q) == 0) {
      // check if the packet queue is to be destroyed
      if (tq->destroy != 0) {
        pthread_mutex_unlock(&(tq->lock));
        pthread_exit(NULL);
      }
      // wait for packet to arrive
      pthread_cond_wait(&(tq->cv), &(tq->lock));
    }

    // get a packet from packet queue
    pkt = packet_q_remove(tq->q);
    assert(pkt);

    pthread_mutex_unlock(&(tq->lock));

    // process packet
    pkt_processor(pkt);

    // done processing, release packet
    free(pkt);
  }
}
